import numpy as np


column = 12
row = 3
a = np.linspace(1, column*row, column*row).reshape(column, row)

def sample_chunk(data, idx_current, idx_relate, idx_padding=-1):
    """
    从 data 数组中沿第 0 维采样一个数据块。

    Args:
        data (np.ndarray): 要从中采样的 N 维数组。采样沿第 0 维进行。
        idx_current (int): 当前的中心索引（沿第 0 维）。
        idx_relate (list or np.ndarray): 相对于 idx_current 的索引列表。
        idx_padding (int): 当采样索引超出 data 第 0 维边界时，
                            用于填充的 idx_relate 中的索引。
                            默认为 -1，表示使用 idx_relate 中的最后一个元素。

    Returns:
        np.ndarray: 采样得到的数据块。形状为 (len(idx_relate),) + data.shape[1:]。
    """
    idx_relate = np.array(idx_relate)
    absolute_indices = idx_current + idx_relate
    # 获取第 0 维的长度
    data_len = data.shape[0]

    # 计算用于填充的绝对索引（沿第 0 维）
    padding_absolute_idx = idx_current + idx_relate[idx_padding]
    # 将填充索引限制在 data 第 0 维的有效范围内
    padding_absolute_idx_clamped = np.clip(padding_absolute_idx, 0, data_len - 1)
    # 获取填充值（沿第 0 维索引，保留其他维度）
    padding_value_slice = data[padding_absolute_idx_clamped, ...]

    # 计算结果数组的形状
    result_shape = (len(absolute_indices),) + data.shape[1:]
    # 创建结果数组，并用填充值切片初始化
    # 注意：这里不能直接用 fill_value=padding_value_slice，因为 fill_value 期望标量
    # 我们先创建一个空数组，然后用广播填充
    sampled_chunk = np.empty(result_shape, dtype=data.dtype)
    sampled_chunk[:] = padding_value_slice # 使用广播填充

    # 找出有效的索引（在 data 第 0 维边界内）
    valid_mask = (absolute_indices >= 0) & (absolute_indices < data_len)

    # 获取有效的绝对索引（沿第 0 维）
    valid_absolute_indices = absolute_indices[valid_mask]

    # 从 data 中获取有效索引对应的值（沿第 0 维索引，保留其他维度）
    valid_values = data[valid_absolute_indices, ...]

    # 将有效值放入结果数组的相应位置（沿第 0 维赋值）
    sampled_chunk[valid_mask, ...] = valid_values

    return sampled_chunk

idx_current_example = 4
idx_relate_example = [-3, 0]
idx_padding_example = 0 # 使用 idx_relate[-1] (即相对索引2) 对应的值来填充

data_example = a.copy()
chunk_2 = sample_chunk(data_example, idx_current_example, idx_relate_example, idx_padding_example)
print(f"Data:\n{data_example}")
print(f"Sampled chunk 2:\n{chunk_2}")


print('end')